//
//  main.cpp
//  germ
//
//  Created by Vinicio Reynoso on 4/3/2010.
//  Copyright (c) 2012 Vinicio Reynoso. All rights reserved.
//

#include "Utility.h"
#include "TabFile.h"
#include "ReadMapper.h"
#include <math.h>
#include <queue>
#include <set>
#include <stdint.h>

using namespace std;

typedef uint64_t u64;
typedef vector<u64> v64;
typedef pair<v64, v64> ReadSet;

typedef vector<bool> vBool;
typedef pair<vBool, vBool> ReadSetBools;	// First vbool records which reads have
// mapped, while the second records which
// reads have mapped to any one query
// sequence.
vector<ReadSet*>      pReadSets;
vector<ReadSetBools*> pReadSetBools;

vector<vector<string> > readSet_out_files;	// Output files for each read set
vector<string> mapping_files;	// Mapping files for each read set

vector<long int> readSetSizes;	// The size of each read set
vector<long int> totalMapped;	// Total number of mapped reads for each set
vector<int>      readLengths;   // The read length for each set

string accessionList = "";
string seqDir = "";
string readSetDir = "";
int    accession_col = -1;
double pID = 0.9;

// function declarations
void usage(void);
bool load_read_set(string& file);
void map_read_sets(ReadMapper& readMapper, bool close);
void record_mapped_reads(int readset);
long int count_mapped_reads(vBool& x);


//****************************************************************************************
//*******************************   FUNCTIONS   ******************************************
//****************************************************************************************
void usage(void)
{
	cout << endl;
	cout << "------------------------- germ2 ----------------------------" << endl;
	cout << "*Purpose: Map given read sets to all currently sequenced" << endl;
	cout << "		   genomes and plasmids." << endl;
	cout << "NOTE: Equal length is assumed for all reads in a set." << endl;
	cout << "	   Genome sequences must be in FASTA format." << endl;
	cout << "Required command line switches:" << endl;
	cout << "	-s	[File listing all read sets with read lengths to be mapped]" << endl;
	cout << "	-g	[File listing gene bank accessions of reference sequences]" << endl;
	cout << "	-c	[Map close matches. true or false]" << endl;
	cout << "	-q	[Directory for annotation and reference sequence files]" << endl;
	cout << "	-r	[Directory for read set files]" << endl;
	cout << "Optional: " << endl;
	cout << "	-P	[Prefix size: default= 20]" << endl;
	cout << "	-I	[Percent Identity for close matches: default= 0.9]" << endl;
	cout << endl;
}

// Used to load a read set file, assumes the ry and ws are in columns 0 and 1
// Returns false if read set could not be loaded.
bool load_read_set (string& file)
{
	TabFile read_setTF;
	if (!read_setTF.open(file, false)) {
		cerr << endl << "IN: main::load_read_set(string& file)" << endl;
		cerr << "\tERROR: could not open read set= " << file << endl;
		return false;
	} else {
		ReadSet * readSet = new pair<v64,v64>( v64(), v64() );
		ReadSetBools * readSetBools = new pair<vBool, vBool>( vBool(), vBool() );
    
    // Get the number of reads.
		unsigned long n_reads = read_setTF.nRows();
		read_setTF.fp_data_start();
    
    // Reserve space for the reads (assumes fastq format)
    n_reads /= 4;
		readSet->first.reserve(n_reads); readSet->second.reserve(n_reads);
		readSetBools->first.reserve(n_reads); readSetBools->second.reserve(n_reads);
		readSetSizes.push_back(n_reads);
    
    // Load the reads.
		u64 r_ry, r_ws;
		vector<string> buff;
		while (read_setTF.getField(buff)) { // gets the read id
			if (buff.empty()) continue;
      
      read_setTF.getField(buff); // gets the read sequence
      
      string read_seq = buff[0];
      int read_length = read_seq.size();
      bool found_bad_char = false;
      r_ry = r_ws = 0;
      
      for (int i = 0; i < read_length; ++i) {
        
        // Add nucleotide to unsigned long ints (r_ry, r_ws)
        switch (read_seq[i]) {
          case 'A' :										        //		 ry	  ws
            r_ry <<= 1; // Add "A" (00)					//	A	 0	  0
            r_ws <<= 1;									        //  G	 0	  1
            break;										          //	T	 1	  0
          case 'G' :                            //	C	 1	  1
            r_ry <<= 1;		// Add "G" (01)
            r_ws <<= 1;	r_ws += 1;
            break;
            
          case 'T' :
            r_ry <<= 1;	r_ry += 1;	// Add "T" (10)
            r_ws <<= 1;
            break;
            
          case 'C' :	// Add "C" (11)
            r_ry <<= 1;	r_ry += 1;
            r_ws <<= 1;	r_ws += 1;
            break;
            
          default :
            found_bad_char = true;
            break;
        }
        
        if (found_bad_char)
          break;
      }
      
      if (found_bad_char == false) {
        readSet->first.push_back(r_ry);
        readSet->second.push_back(r_ws);
        
        readSetBools->first.push_back(false);
        readSetBools->second.push_back(false);
      }
      
      
      read_setTF.getField(buff); // gets the description
      read_setTF.getField(buff); // gets the quality values
      
		}
		
		// Push back the read set.
		pReadSets.push_back(readSet);
		pReadSetBools.push_back(readSetBools);
		
		read_setTF.close();
	}
	return true;
}
// Maps all currently loaded read sets.
void map_read_sets (ReadMapper& read_mapper, bool close)
{
	TabFile accessionsTF;
	if (!accessionsTF.open(accessionList)) {
		cerr << endl << "IN: main::map_read_sets(...);" << endl;
		cerr << "\tERROR: could not open accessions file= " << endl;
		return;
	}
	else {
		vector<string> buff;
		while (accessionsTF.getField(buff)) {
			if (buff.empty()) continue;
			
			string temp = buff[accession_col];	// Looks like 'accession.version_num'
			string accession = temp.substr(0, temp.find_first_of('.'));
			string fna_file = accession + ".fna";
			
			int kmersize = readLengths[0];
			
			read_mapper.load_seq(kmersize, fna_file, seqDir, true);
			
			int n_sets = pReadSets.size();
			for (int i = 0; i < n_sets; ++i) {
				
				// Check that the sequence was decomposed into readlength k-mers.
				if (readLengths[i] != kmersize) {
					kmersize = readLengths[i];
					read_mapper.load_seq(kmersize, fna_file, seqDir, true);
				}
				
				// Map the read set, not allowing mismatches.
				read_mapper.map_reads(pReadSets[i]->first, pReadSets[i]->second,
                              pReadSetBools[i]->second);
				
				// Map the read set, allowing mismatches.
				if (close) {
					int allowed_mm = kmersize * pID;
					int threshold = kmersize - allowed_mm;
					
					read_mapper.map_close_reads(pReadSets[i]->first,
                                      pReadSets[i]->second,
                                      pReadSetBools[i]->second, threshold);
				}
				
				record_mapped_reads(i);
				
				read_mapper.filter_mappings();
				
				// Write mappings to output files
				read_mapper.write_mappings  (readSet_out_files[i][0], true);
				read_mapper.write_gene_hits (readSet_out_files[i][1]);
				read_mapper.write_seq_hits  (readSet_out_files[i][2]);
				
				// Write summary to screen.
				cout << '\t' << read_mapper.total_hits() << '\t';
				cout << read_mapper.n_close_hits() << '\t';
				cout << read_mapper.n_exact_hits() << '\t';
				cout << read_mapper.n_eNC_hits() << '\t';
				cout << read_mapper.n_cNC_hits() << '\t';
				cout << read_mapper.n_eP_hits() << '\t';
				cout << read_mapper.n_cP_hits() << '\t';
				cout << read_mapper.n_eR_hits() << '\t';
				cout << read_mapper.n_cR_hits() << '\t';
				cout << read_mapper.accession() << '\t';
				
				int seconds = read_mapper.secs_spent_mapping();
				double mins = seconds/60.0;
				cout << fixed << setprecision(2) << mins << '\t';
				
				string temp = readSet_out_files[i][0];
				size_t pos = temp.find('_', temp.find_first_of('_')+1);
				temp = temp.substr(pos+1);
				temp.erase(temp.find_last_of('.'));
				cout << temp << '\n';
				
				// Reset counts and mappings for the next read set.
				read_mapper.reset_counts();
			}
		}
	}
}
// Sets each read's bool value to true if it mapped onto a sequence.
// Bool values in the second vector will be set to back to false.
void record_mapped_reads (int readSet)
{
	long int n_reads = pReadSetBools[readSet]->first.size();
	for (long int i = 0; i < n_reads; ++i) {
    
		if (pReadSetBools[readSet]->second.at(i) == true) {
			pReadSetBools[readSet]->first.at(i)  = true;
			pReadSetBools[readSet]->second.at(i) = false;
		}
    
	}
}
// Counts the number of mapped reads and returns the number of reads.
// The first vBool in the pair should be passed, not the second.
long int count_mapped_reads (int readSet)
{
	long int stop = pReadSetBools[readSet]->first.size();
	long int n_mapped = 0;
	
	for (long int i = 0; i < stop; ++i)
		if(pReadSetBools[readSet]->first.at(i) == true)
			++n_mapped;
	
	return n_mapped;
}
//****************************************************************************************
//***********************************  MAIN  *********************************************
//****************************************************************************************
int main (int argc, char* argv[]){
  //get command line values and initialize some variables
	Utility io;
	if (io.getArguments(argc, argv) == false) {
		cerr << endl << "Error initializing Utility object." << endl;
		cerr << "Odd number of arguments passed on comand line." << endl << endl;
		usage();
		return 0;
	}
	if (io.argvSize() == 0) {
	 	cerr << endl << "ERROR: required command line values missing" << endl;
		usage();
		return 0;
	}
	int prefix_size = 20;
	bool useClose;
	bool map_one_set_at_a_time = false;
	string set_list_file = "";
  //map switiches
	if (!io.mapSwitch("-s", set_list_file)) {
		cerr << "\n** ERROR: read set list missing." << endl; usage(); return 0;
	}
	
	if (!io.mapSwitch("-g", accessionList)) {
		cerr << "\n** ERROR: accessions file missing." << endl; usage(); return 0;
	}
	
	if (!io.mapSwitch("-c", useClose)) {
		cerr << "\n** ERROR: Option for close wasn't given." << endl; usage(); return 0;
	}
	
	if (!io.mapSwitch("-q", seqDir)) {
		cerr << "\n** ERROR: No seq and annotation dir." << endl; usage(); return 0;
	}
	
	if (!io.mapSwitch("-r", readSetDir)) {
		cerr << "\n** ERROR: No dir for read set files." << endl; usage(); return 0;
	}
	
	io.mapSwitch("-P", prefix_size);
	io.mapSwitch("-I", pID);
  
	TabFile set_listTF, accessionsTF;
	vector<string> buff;
	
	// Get the accession number column
	if (!accessionsTF.open(accessionList)) {
		cerr << "IN: main()" << endl;
		cerr << "\tERROR: could not open input file= " << accessionList << endl;
		return 0;
	}
	
	accession_col = accessionsTF.getFieldIndex("* Which field contains: accessions ");
	accessionsTF.close();
  
  // Make output file for each read set and store read lengths.
	if (!set_listTF.open(set_list_file, false)) {
		cerr << "IN: main()" << endl;
		cerr <<	"\tERROR: could not open the set list file= " << set_list_file << endl;
		return 0;
	}
	else {
		while (set_listTF.getField(buff))
		{
			if (buff.empty()) continue;
			
			string readFile = buff[0];
			readLengths.push_back(atoi(buff[1].c_str()));
      
			string temp = readFile.substr(0, readFile.find_last_of('.'));
			string read_mappings_file = "read_mappings_" + temp + ".txt";
			string gene_hits_file = "gene_hits_" + temp + ".txt";
			string seq_hits_file = "seq_hits_" + temp + ".txt";
			
			// Print the headers to output files.
			ofstream out;
			out.open(read_mappings_file.c_str());
			if (!out) {
				cerr << "ERROR: could not create output file= ";
				cerr << read_mappings_file << endl;
				return 0;
			}
			out << "ry_read\tid_read\tmm\tloc\tgene_syn\tgene_strand\taccession\tseq";
			out << endl;
			out.close();
			
			
			out.open(gene_hits_file.c_str());
			if (!out) {
				cerr << "ERROR: could not create output file= ";
				cerr << gene_hits_file << endl;
				return 0;
			}
			out << "seq\taccession\tp=0,r=1\tgene_name\tgene_syn\tstrand\t";
			out << "eHits\tcHits\tproduct" << endl;
			out.close();
			
			
			out.open(seq_hits_file.c_str());
			if (!out) {
				cerr << "ERROR: could not open output file= " << seq_hits_file << endl;
				return 0;
			}
			out << "seq\taccession\ttotal\tclose\tfclose\texact\tfexact\t";
			out << "eProt\tcProt\teRna\tcRna\teNCoding\tcNCoding" << endl;
			out.close();
			
			vector<string> output_files;
			output_files.push_back(read_mappings_file);
			output_files.push_back(gene_hits_file);
			output_files.push_back(seq_hits_file);
			
			readSet_out_files.push_back(output_files);
		}
	}
	// Move to the first read set in the list
	set_listTF.fp_data_start();
	
	// Create a read mapper object
	ReadMapper read_mapper(prefix_size);
	
	// Load the read sets.
	while (set_listTF.getField(buff)) {
		if (buff.empty()) continue;
    
		cout << "\nLoading read set file= " << buff[0] << flush;
		if (readSetDir == "") {
			if ( !load_read_set(buff[0]) )
				continue;
		} else {
			string temp = readSetDir + buff[0];
			if ( !load_read_set(temp) )
				continue;
		}
		cout << "\tDone!";
		
		// Map the read sets.
		if (map_one_set_at_a_time) {
      
			cout << "\n\n\tT\tC\tE\teNC\tcNC\teProt\tcProt\teRNA\t";
			cout << "cRNA\tSeq\t\tMins\tReadSet" << endl;
      
			map_read_sets(read_mapper, useClose);
			
			// Record number of reads and the number of mapped reads.
			long int n_mapped = count_mapped_reads(0);
			totalMapped.push_back(n_mapped);
			
			mapping_files.push_back(readSet_out_files[0][0]);
			
			readSet_out_files.erase(readSet_out_files.begin());
			readLengths.erase(readLengths.begin());
			
			// Clear read set data.
			delete pReadSets[0];
			pReadSets.clear();
			
			delete pReadSetBools[0];
			pReadSetBools.clear();
		}
	}
	
	if (!map_one_set_at_a_time) {
		// Write summary header to screen.
		cout << "\n\n\tT\tC\tE\teNC\tcNC\teProt\tcProt\teRNA\tcRNA\t";
		cout << "Seq\t\tMins\tReadSet" << endl;
		
		// Map all read sets.
		map_read_sets(read_mapper, useClose);
		
		// Record number of reads and the number of mapped reads.
		int n_sets = pReadSets.size();
		for (int i = 0; i < n_sets; ++i) {
			long int n_mapped = count_mapped_reads(i);
			
			totalMapped.push_back(n_mapped);
			mapping_files.push_back(readSet_out_files[i][0]);
			
			// Clear memory for read sets.
			delete pReadSets[i];
			delete pReadSetBools[i];
		}
	}
	cout << endl << endl;
	
	// Filter the aggregate read mappings for each read set
	int n_mapping_files = mapping_files.size();
	for (int i = 0; i < n_mapping_files; ++i) {
		read_mapper.filter_mappings(mapping_files[i]);
		
		// Write the total number of reads and mapped reads for each set.
		string temp = mapping_files[i];
		size_t pos = temp.find('_', temp.find_first_of('_')+1);
		temp = temp.substr(pos+1);
		temp.erase(temp.find_last_of('.'));
		
		long double prct_mapped = ((long double)totalMapped[i] /
                               (long double)readSetSizes[i]) * 100.00;
		
		cout << temp << ": ";
		cout << fixed << setprecision(1);
		cout << prct_mapped << "% of ";
		
		stringstream ss;
		ss << readSetSizes[i];
		string n_reads_str = ss.str();
		int len = n_reads_str.size() - 1;
		int digit = 1;
		
		for (int j = len; j >= 0; --j) {
			if ( digit%3 == 0 && j != 0)
				n_reads_str.insert(j, ",");
			++digit;
		}
		cout << n_reads_str << " reads mapped" << endl << endl;
	} 
	return 0;
}
